from sentence_transformers import SentenceTransformer
import chromadb
from llm.doubao import doubao_qa
from llm.local import ollama_qa

model = SentenceTransformer("all-MiniLM-L6-v2")

client = chromadb.PersistentClient(path="./chromadb_data")
collection = client.get_or_create_collection("rag")


def get_query_embedding(query):
    return model.encode(query).tolist()


def retrieve_related_chunks(query_embedding, n_results=3):
    results = collection.query(query_embeddings=[query_embedding], n_results=n_results)
    related_chunks = results.get("documents")
    if not related_chunks or not related_chunks[0]:
        exit(1)
    return related_chunks[0]


# 槽位状态管理类
# {
#  "skill": "Java",
#  "experience": 3,
#  "salary": 20000,
#  "remote": "是"
# }
class SlotState:
    def __init__(self):
        self.slots = {}

    def update(self, key, value):
        self.slots[key] = value

    def remove(self, key):
        if key in self.slots:
            del self.slots[key]

    def to_sql_where(self):
        conds = []
        for k, v in self.slots.items():
            if k == "experience":
                conds.append(f"experience>={v}")
            if k == "salary":
                conds.append(f"salary <= {v}")
            if k == "remote":
                conds.append(f"remote = '{v}'")
            if k == "skill":
                conds.append(f"skill LIKE '%{v}%'")
        return " AND ".join(conds)


def parse_slots(user_input):
    prompt = """
     请从下面的招聘筛选需求中抽取结构化的槽位，输出JSON格式
     支持字段：skill experience salary remote
     用户输入{user_input}
     槽位:
    """.format(
        user_input=user_input
    )
    import json

    slots_json = doubao_qa(prompt)
    try:
        start = slots_json.index("{")
        end = slots_json.rindex("}")
        slots_json = slots_json[start : end + 1]
        slots = json.loads(slots_json)
    except Exception:
        slots = {}
    print(f"slots_json：{slots_json}")
    return slots


if __name__ == "__main__":
    # 支持多轮补充条件，自动解析槽位并动态生成查询SQL
    state = SlotState()
    while True:
        user_input = input("请输入招聘的条件(回车结束)")
        if not user_input.strip():
            break
        print(f"user_input:{user_input}")
        slots = parse_slots(user_input)
        for k, v in slots.items():
            state.update(k, v)
        sql_where = state.to_sql_where()
        print(sql_where)
# 三年以上Java开发经验
# 期望薪资低于2万
# 能远程办公
